import cv2
import numpy as np
import pyscreenshot as ImageGrab
import pyautogui
from config import SCREENSHOT_REGION, WINDOW_TITLE

def capture_screen(region=None):
    """
    捕获屏幕截图
    
    Args:
        region (tuple): 截图区域 (x, y, width, height)，如果为None则捕获全屏
    
    Returns:
        numpy.ndarray: OpenCV格式的图像数据
    """
    if region is None:
        # 捕获全屏
        screenshot = pyautogui.screenshot()
    else:
        # 捕获指定区域
        screenshot = pyautogui.screenshot(region=region)
    
    # 转换为OpenCV格式
    open_cv_image = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
    return open_cv_image

def capture_window(window_title):
    """
    捕获指定窗口（此功能可能需要平台特定的实现）
    
    Args:
        window_title (str): 窗口标题
    
    Returns:
        numpy.ndarray: OpenCV格式的图像数据
    """
    # 注意：精确捕获特定窗口在不同操作系统上实现方式不同
    # 这里使用一个简化的方法，通过查找窗口位置来确定截图区域
    try:
        # 查找窗口位置
        window_location = pyautogui.locateOnScreen(window_title + ".png", confidence=0.9)
        if window_location:
            # 如果找到了窗口的标识图片，则以此为基准截图
            screenshot = pyautogui.screenshot(region=window_location)
            open_cv_image = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
            return open_cv_image
        else:
            # 如果没找到，则回退到普通屏幕捕获
            return capture_screen(SCREENSHOT_REGION)
    except:
        # 出现异常时回退到普通屏幕捕获
        return capture_screen(SCREENSHOT_REGION)

def main():
    # 测试屏幕捕获功能
    image = capture_screen(SCREENSHOT_REGION)
    cv2.imwrite("test_screenshot.png", image)
    print("屏幕截图已保存为 test_screenshot.png")

if __name__ == "__main__":
    main()